#!/usr/bin/env python3
"""
医学图像分割推理脚本

使用方法:
    python inference.py
    python inference.py --config custom_config.yaml
    python inference.py --mode single --input_path /path/to/your/image.png
    python inference.py --mode batch --input_path /path/to/your/image_folder
    python inference.py --mode batch --input_path ./data/test.txt
    python inference.py --model_path ./saved_models/best_model.pth
    python inference.py --no_visualize
"""

import argparse
import yaml
import os
import sys
import torch
from tqdm import tqdm
from PIL import Image
import numpy as np

from data.data_loader import create_single_image_loader, create_inference_loader, get_num_classes
from utils.model_utils import load_model_for_inference
from utils.plot_utils import visualize_prediction, save_prediction_overlay


def parse_args():
    """解析命令行参数"""
    
    parser = argparse.ArgumentParser(description="医学图像分割推理脚本")
    parser.add_argument("--config", type=str, default="config.yaml",
                       help="配置文件路径 (默认: config.yaml)")
    parser.add_argument("--mode", type=str, choices=["single", "batch"], default=None,
                       help="推理模式: single (单张图像) 或 batch (批量图像) (默认: 使用config中的设置)")
    parser.add_argument("--input_path", type=str, default=None,
                       help="输入路径 (默认: 使用config中的设置)")
    parser.add_argument("--model_path", type=str, default=None,
                       help="模型权重路径 (默认: 使用config中的路径)")
    parser.add_argument("--no_visualize", action="store_true",
                       help="不保存可视化结果")
    
    return parser.parse_args()


def load_config(config_path):
    """加载配置文件"""
    
    if not os.path.exists(config_path):
        print(f"错误: 配置文件不存在: {config_path}")
        sys.exit(1)
    
    try:
        with open(config_path, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)
        return config
    except Exception as e:
        print(f"错误: 无法加载配置文件: {e}")
        sys.exit(1)


def main():
    """主函数"""
    
    # 解析命令行参数
    args = parse_args()
    
    # 加载配置
    config = load_config(args.config)
    
    # 覆盖config中的推理设置（如果命令行参数指定）
    inference_mode = args.mode if args.mode else config["inference"]["mode"]
    input_list_path = args.input_path if args.input_path else config["inference"]["input_list"]
    image_dir = config["inference"]["image_dir"]
    model_path = args.model_path if args.model_path else config["inference"]["model_path"]
    visualize_results = not args.no_visualize and config["inference"]["visualize"]
    output_dir = config["inference"]["output_dir"]
    
    # 获取类别颜色和名称
    class_colors = config["inference"]["colors"]
    num_classes = get_num_classes(config)
    class_names = config["data"]["name_classes"] if "name_classes" in config["data"] else [f"Class {i}" for i in range(num_classes)]

    # 检查模型路径
    if not model_path or not os.path.exists(model_path):
        print(f"错误: 模型文件不存在或未指定: {model_path}. 请在config.yaml中设置inference.model_path或通过--model_path参数指定。")
        sys.exit(1)
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(os.path.join(output_dir, "masks"), exist_ok=True)
    if visualize_results:
        os.makedirs(os.path.join(output_dir, "visualizations"), exist_ok=True)

    print(f"开始推理，模式: {inference_mode}")
    
    try:
        # 加载模型
        model, device = load_model_for_inference(model_path, config)
        if model is None:
            print("模型加载失败")
            sys.exit(1)
        
        if inference_mode == "single":
            # 单图模式：从命令行输入图像路径
            if not input_list_path or input_list_path.endswith('.txt'):
                # 如果未指定路径或路径是txt文件，提示用户输入
                input_list_path = input("请输入单张图像路径: ").strip()
            
            # 检查文件是否存在且是图像文件
            if not os.path.exists(input_list_path):
                print(f"错误: 文件不存在: {input_list_path}")
                sys.exit(1)
            if not input_list_path.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.gif')):
                print(f"错误: 不是有效的图像文件: {input_list_path}")
                sys.exit(1)
                
            image_loader = create_single_image_loader(config, input_list_path)
            original_image = Image.open(input_list_path).convert("RGB")
            file_name = os.path.splitext(os.path.basename(input_list_path))[0]
            
            with torch.no_grad():
                for images in tqdm(image_loader, desc="推理中"):
                    images = images.to(device)
                    outputs = model(images)
                    pred_mask = torch.argmax(outputs, dim=1).squeeze(0).cpu().numpy()
                    
                    # 保存预测掩码
                    mask_save_path = os.path.join(output_dir, "masks", f"{file_name}_mask.png")
                    Image.fromarray(pred_mask.astype(np.uint8) * 255).save(mask_save_path)
                    print(f"预测掩码已保存到: {mask_save_path}")
                    
                    # 保存可视化结果
                    if visualize_results:
                        vis_save_path = os.path.join(output_dir, "visualizations", f"{file_name}_overlay.png")
                        save_prediction_overlay(original_image, pred_mask, vis_save_path, class_colors_list=class_colors)
                        print(f"可视化结果已保存到: {vis_save_path}")

        elif inference_mode == "batch":
            # 批量模式：从input_list读取文件名列表
            if not os.path.exists(input_list_path):
                print(f"错误: 文件列表不存在: {input_list_path}")
                sys.exit(1)
            
            # 读取文件名列表（不带后缀）
            with open(input_list_path, "r") as f:
                file_names = [line.strip() for line in f.readlines()]
            
            if not file_names:
                print(f"警告: 未在 {input_list_path} 中找到任何图像文件")
                sys.exit(0)
            
            # 创建数据加载器
            image_loader, _ = create_inference_loader(
                config=config,
                input_path=input_list_path,
                batch_size=1, # config["training"]["batch_size"]
            )
            
            # 使用config中的image_dir作为图像根目录
            image_dir = config["inference"]["image_dir"]
            
            for i, (images, _) in enumerate(tqdm(image_loader, desc="批量推理中")):
                images = images.to(device)  # 这里应该是将图像张量移动到设备上
                outputs = model(images)
                pred_masks = torch.argmax(outputs, dim=1).cpu().numpy()
                
                for j, pred_mask in enumerate(pred_masks):
                    current_file_name = file_names[i * image_loader.batch_size + j]
                    original_image_path = os.path.join(image_dir, current_file_name + ".png")
                    original_image = Image.open(original_image_path).convert("RGB")

                    # 保存预测掩码
                    mask_save_path = os.path.join(output_dir, "masks", f"{current_file_name}_mask.png")
                    Image.fromarray(pred_mask.astype(np.uint8) * 255).save(mask_save_path)
                    
                    # 保存可视化结果
                    if visualize_results:
                        vis_save_path = os.path.join(output_dir, "visualizations", f"{current_file_name}_overlay.png")
                        save_prediction_overlay(original_image, pred_mask, vis_save_path, class_colors_list=class_colors)
            print(f"批量推理完成！结果已保存到 {output_dir}/masks 和 {output_dir}/visualizations 目录")

        else:
            print(f"错误: 不支持的推理模式: {inference_mode}")
            sys.exit(1)
            
        print("推理完成！")
        
    except Exception as e:
        print(f"推理过程中发生错误: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()
